package site.wanjiahao.tree;

/**
 * AVL树是一个比较特殊的二叉搜索树
 * AVL树的每个节点的平衡因子不超过1
 */
public class AVLTree<E> extends BST<E> {

    /**
     * 添加节点方法后，模板方法设计模式
     *
     * @param node node
     */
    @Override
    protected void addAfter(Node node) {
        AVLNode avlNode = toAVLNode(node);
        // 找到当前节点的祖先节点的失衡节点
        while ((avlNode = toAVLNode(avlNode.parent)) != null) {
            if (avlNode.isBalance()) {
                // 当前节点平衡，计算当前节点的高度，给下一次计算使用
                avlNode.updateHeight();
            } else {
                // 不平衡，开始进行旋转操作，旋转完成之后，整颗数一定处于平衡状态
                avlNode.reBalance();
                break;
            }
        }
    }

    @Override
    protected void removeAfter(Node node, Node replacement) {
        AVLNode avlNode = toAVLNode(node);
        // 找到当前节点的祖先节点的失衡节点
        while ((avlNode = toAVLNode(avlNode.parent)) != null) {
            if (avlNode.isBalance()) {
                // 当前节点平衡，计算当前节点的高度，给下一次计算使用
                avlNode.updateHeight();
            } else {
                // 不平衡，开始进行旋转操作，旋转完成之后，不一定处于平衡，需要再次计算
                avlNode.reBalance();
            }
        }
    }

    @Override
    protected Node createNode(E ele, Node parent) {
        return new AVLNode(ele, parent);
    }

    public AVLNode toAVLNode(Node node) {
        return (AVLNode) node;
    }

    private class AVLNode extends Node {
        // 多增加属性 height
        int height = 1;

        public AVLNode(E ele, Node parent) {
            super(ele, parent);
        }

        // 判断当前节点是否失衡
        public boolean isBalance() {
            int leftHeight = left == null ? 0 : toAVLNode(left).height;
            int rightHeight = right == null ? 0 : toAVLNode(right).height;
            return Math.abs(leftHeight - rightHeight) <= 1;
        }

        // 更新当前高度
        public void updateHeight() {
            height = getChildMaxHeight() + 1;
        }

        // 重新恢复平衡
        public void reBalance() {
            // 寻找最大高度的子节点，孙节点
            AVLNode son = tallerChild();
            AVLNode grandSon = son.tallerChild();
            // 判断旋转类型
            if (son.isLeftNode()) {
               /* if (grandSon.isRightNode()) {
                    // LR
                    son.rotateLeft();
                }
                rotateRight();*/
                if (grandSon.isLeftNode()) {
                    // LL
                    rotate(grandSon, grandSon.right, son, son.right, this);
                } else {
                    // LR
                    rotate(son, grandSon.left, grandSon, grandSon.right, this);
                }
            } else {
               /* if (grandSon.isLeftNode()) {
                    // RL
                    son.rotateRight();
                }
                rotateLeft();*/
                if (grandSon.isLeftNode()) {
                    // RL
                    rotate(this, grandSon.left, grandSon, grandSon.right, son);
                } else {
                    // RR
                    rotate(this, son.left, son, grandSon.left, grandSon);
                }
            }
        }

        @Override
        public void rotateAfter(Node cur, Node next, Node node) {
            super.rotateAfter(cur, next, node);

            // 更新高度变化
            toAVLNode(cur).updateHeight();
            toAVLNode(next).updateHeight();
        }

        @Override
        public void rotate(Node b, Node c, Node d, Node e, Node f) {
            super.rotate(b, c, d, e, f);

            // 更新高度
            toAVLNode(b).updateHeight();
            toAVLNode(f).updateHeight();
            toAVLNode(d).updateHeight();
        }

        // 寻找最大高度的子节点
        public AVLNode tallerChild() {
            int leftHeight = getLeftHeight();
            int rightHeight = getRightHeight();
            if (leftHeight > rightHeight) return toAVLNode(left);
            else if (leftHeight < rightHeight) return toAVLNode(right);
            else if (this.isLeftNode()) return toAVLNode(left);
            else return toAVLNode(right);
        }

        public int getLeftHeight() {
            return left == null ? 0 : toAVLNode(left).height;
        }

        public int getRightHeight() {
            return right == null ? 0 : toAVLNode(right).height;
        }

        public int getChildMaxHeight() {
            return Math.max(getLeftHeight(), getRightHeight());
        }

        @Override
        public String toString() {
            return ele + "";
        }
    }

}
